{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "\n",
    "RATE = 8\n",
    "A = 12\n",
    "B = 6\n",
    "\n",
    "def bytes_to_int(bytes):\n",
    "    return sum((b << (8 * i) for i, b in enumerate(reversed(bytes))))\n",
    "\n",
    "def int_to_bytes(value, length):\n",
    "    return bytes((value >> (8 * i) & 0xff for i in range(length - 1, -1, -1)))\n",
    "\n",
    "def bytes_to_state(bytes):\n",
    "    return [bytes_to_int(bytes[i * 8:(i + 1) * 8]) for i in range(5)]\n",
    "\n",
    "def rotr(x, n):\n",
    "    return (x >> n) | (x << (64 - n)) & 0xFFFFFFFFFFFFFFFF\n",
    "\n",
    "def ascon_permutation(S, rounds):\n",
    "    for r in range(12 - rounds, 12):\n",
    "        S[2] ^= 0xf0 - r * 0x10 + r * 0x1\n",
    "        S[0] ^= S[4]\n",
    "        S[4] ^= S[3]\n",
    "        S[2] ^= S[1]\n",
    "        T = [(S[i] ^ 0xFFFFFFFFFFFFFFFF) & S[(i + 1) % 5] for i in range(5)]\n",
    "        for i in range(5):\n",
    "            S[i] ^= T[(i + 1) % 5]\n",
    "        S[1] ^= S[0]\n",
    "        S[0] ^= S[4]\n",
    "        S[3] ^= S[2]\n",
    "        S[2] ^= 0xFFFFFFFFFFFFFFFF\n",
    "        S[0] ^= rotr(S[0], 19) ^ rotr(S[0], 28)\n",
    "        S[1] ^= rotr(S[1], 61) ^ rotr(S[1], 39)\n",
    "        S[2] ^= rotr(S[2], 1) ^ rotr(S[2], 6)\n",
    "        S[3] ^= rotr(S[3], 10) ^ rotr(S[3], 17)\n",
    "        S[4] ^= rotr(S[4], 7) ^ rotr(S[4], 41)\n",
    "\n",
    "def ascon_initialize(S, key, nonce):\n",
    "    k = len(key) * 8\n",
    "    iv_zero_key_nonce = bytes([k, RATE * 8, A, B]) + (20-len(key)) * b\"\\x00\" + key + nonce\n",
    "    S[0], S[1], S[2], S[3], S[4] = bytes_to_state(iv_zero_key_nonce)\n",
    "    ascon_permutation(S, A)\n",
    "    zero_key = bytes_to_state((40-len(key)) * b\"\\x00\" + key)\n",
    "    for i in range(5):\n",
    "        S[i] ^= zero_key[i]\n",
    "\n",
    "def ascon_process_associated_data(S, associated_data):\n",
    "    if associated_data:\n",
    "        a_padding = b\"\\x80\" + (RATE - (len(associated_data) % RATE) - 1) * b\"\\x00\"\n",
    "        a_padded = associated_data + a_padding\n",
    "        for i in range(0, len(a_padded), RATE):\n",
    "            S[0] ^= bytes_to_int(a_padded[i:i+8])\n",
    "            if RATE == 16:\n",
    "                S[1] ^= bytes_to_int(a_padded[i+8:i+16])\n",
    "            ascon_permutation(S, B)\n",
    "    S[4] ^= 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ascon_process_plaintext(S, plaintext):\n",
    "    p_lastlen = len(plaintext) % RATE\n",
    "    p_padding = b\"\\x80\" + (RATE - p_lastlen - 1) * b\"\\x00\"\n",
    "    p_padded = plaintext + p_padding\n",
    "    ciphertext = bytes()\n",
    "    for i in range(0, len(p_padded) - RATE, RATE):\n",
    "        S[0] ^= bytes_to_int(p_padded[i:i+8])\n",
    "        ciphertext += int_to_bytes(S[0], 8)\n",
    "        ascon_permutation(S, B)\n",
    "    return ciphertext + int_to_bytes(S[0], 8)[:p_lastlen]\n",
    "\n",
    "def ascon_process_ciphertext(S, ciphertext):\n",
    "    c_lastlen = len(ciphertext) % RATE\n",
    "    c_padded = ciphertext + (RATE - c_lastlen) * b\"\\x00\"\n",
    "    plaintext = bytes()\n",
    "    for i in range(0, len(c_padded) - RATE, RATE):\n",
    "        Ci = bytes_to_int(c_padded[i:i+8])\n",
    "        plaintext += int_to_bytes(S[0] ^ Ci, 8)\n",
    "        S[0] = Ci\n",
    "        ascon_permutation(S, B)\n",
    "    return plaintext + int_to_bytes(S[0], 8)[:c_lastlen]\n",
    "\n",
    "def ascon_finalize(S, key):\n",
    "    S[1] ^= bytes_to_int(key[0:8])\n",
    "    S[2] ^= bytes_to_int(key[8:16])\n",
    "    ascon_permutation(S, A)\n",
    "    S[3] ^= bytes_to_int(key[0:8])\n",
    "    S[4] ^= bytes_to_int(key[8:16])\n",
    "    return int_to_bytes(S[3], 8) + int_to_bytes(S[4], 8)\n",
    "\n",
    "def ascon_aead_encrypt(key, nonce, associated_data, plaintext):\n",
    "    assert len(key) == 16 and len(nonce) == 16\n",
    "    S = [0] * 5\n",
    "    ascon_initialize(S, key, nonce)\n",
    "    ascon_process_associated_data(S, associated_data)\n",
    "    ciphertext = ascon_process_plaintext(S, plaintext)\n",
    "    tag = ascon_finalize(S, key)\n",
    "    return ciphertext + tag\n",
    "\n",
    "def ascon_aead_decrypt(key, nonce, associated_data, ciphertext):\n",
    "    assert len(key) == 16 and len(nonce) == 16 and len(ciphertext) >= 16\n",
    "    S = [0] * 5\n",
    "    ascon_initialize(S, key, nonce)\n",
    "    ascon_process_associated_data(S, associated_data)\n",
    "    plaintext = ascon_process_ciphertext(S, ciphertext[:-16])\n",
    "    tag = ascon_finalize(S, key)\n",
    "    if tag == ciphertext[-16:]:\n",
    "        return plaintext\n",
    "    else:\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tbody>\n",
       "<tr><td>Key               </td><td>000102030405060708090a0b0c0d0e0f                                </td></tr>\n",
       "<tr><td>Nonce             </td><td>101112131415161718191a1b1c1d1e1f                                </td></tr>\n",
       "<tr><td>Associated Data   </td><td>202122232425262728292a2b2c2d2e2f                                </td></tr>\n",
       "<tr><td>Plaintext         </td><td>303132333435363738393a3b3c3d3e3f                                </td></tr>\n",
       "<tr><td>Ciphertext        </td><td>0ef907dbcd17b7e4f8c9318c8444970defa74b5a49ce0d6ceb46e3a40886f4d4</td></tr>\n",
       "<tr><td>Received Plaintext</td><td>303132333435363738393a3b3c3d3e3f                                </td></tr>\n",
       "</tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "key   = b\"\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07\\x08\\x09\\x0A\\x0B\\x0C\\x0D\\x0E\\x0F\"\n",
    "nonce = b\"\\x10\\x11\\x12\\x13\\x14\\x15\\x16\\x17\\x18\\x19\\x1A\\x1B\\x1C\\x1D\\x1E\\x1F\"\n",
    "associated_data = b\"\\x20\\x21\\x22\\x23\\x24\\x25\\x26\\x27\\x28\\x29\\x2A\\x2B\\x2C\\x2D\\x2E\\x2F\"\n",
    "plaintext      = b\"\\x30\\x31\\x32\\x33\\x34\\x35\\x36\\x37\\x38\\x39\\x3A\\x3B\\x3C\\x3D\\x3E\\x3F\"\n",
    "\n",
    "ciphertext = ascon_aead_encrypt(key, nonce, associated_data, plaintext)\n",
    "received_plaintext = ascon_aead_decrypt(key, nonce, associated_data, ciphertext)\n",
    "\n",
    "if received_plaintext is None:\n",
    "    print(\"Verification failed!\")\n",
    "    \n",
    "else:\n",
    "    # Jupyter table with the results\n",
    "    from IPython.display import HTML, display\n",
    "    import tabulate\n",
    "    table = [[\"Key\", key.hex()],\n",
    "             [\"Nonce\", nonce.hex()],\n",
    "             [\"Associated Data\", associated_data.hex()],\n",
    "             [\"Plaintext\", plaintext.hex()],\n",
    "             [\"Ciphertext\", ciphertext.hex()],\n",
    "             [\"Received Plaintext\", received_plaintext.hex()]]\n",
    "    display(HTML(tabulate.tabulate(table, tablefmt='html')))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}